import itk
import numpy as np
import nibabel as nib
import os
import pandas as pd

from tqdm import tqdm
'''
    为了兼容性考虑，这里python用3.7,然后后面安装其他包的时候不用指定版本，如果3.9会报很多乱七八糟的错
    使用代码前需要通过pip安装如下环境：tqdm（用于进度条展示），nibabel（用于数据读取），itk-elastix（用于运算配准），openpyxl和pandas（用于读取excel批量裁切）其他的包应该环境自带了
'''


def listdir(path):
    """
    用来遍历文件夹下的文件，并且返回文件名字列表
    """
    list_name = []
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        list_name.append(file_path)
    list_name.sort()
    return list_name


def reg(movingimgdir, fixedimgdir, savedir, labeldir=None, savelabeldir=None, overwrite=True, log=False):
    """
    配准函数
    """
    #   检查存放路径是否存在，不存在的话递归创建
    if not os.path.exists(savedir):
        os.makedirs(savedir)

    #   将需要处理的图像名字都读出来，因为对应的图像名字一般是一样的，所有读一个就行了
    file_names = listdir(movingimgdir)

    #   遍历所有的文件名字
    for file_name in tqdm(file_names):
        #   根据需要修改各个路径下文件对应的后缀名，比如raw,nii,dicom之类的
        movingimgpath = os.path.join(movingimgdir, os.path.splitext(os.path.basename(file_name))[0] + '.nii')
        fixed = os.path.join(fixedimgdir, os.path.splitext(os.path.basename(file_name))[0] + '.nii')
        saveimg_ = os.path.join(savedir, os.path.splitext(os.path.basename(file_name))[0] + '.nii')
        if labeldir is not None:
            labelpath = os.path.join(labeldir, os.path.splitext(os.path.basename(file_name))[0] + '.nii')
        if savelabeldir is not None:
            savelbl = os.path.join(savelabeldir, os.path.splitext(os.path.basename(file_name))[0] + '.nii')

        #   如果开启了overwrite模式，则每次跑代码都会全部配准，关闭的话，就只配准还没配的
        if not overwrite:
            if os.path.exists(saveimg_):
                print(saveimg_)
                continue
        try:
            #   根据目标文件的路径，利用nibabel读取文件
            fixedimg = nib.load(fixed)
            affine = fixedimg.affine
            fixedimg = fixedimg.get_fdata()
            fixedimg = fixedimg.astype(np.float32)

            #   根据待配准文件的路径，读取
            movingimg = nib.load(movingimgpath)
            movingimg = movingimg.get_fdata()
            movingimg = movingimg.astype(np.float32)
            #   如果是raw的读取则把上面三行注释掉，调用下面的这几行代码
            # movingimg = np.fromfile(file=movingimgpath, dtype=np.int16)
            # t = movingimg.shape[0] / 512 / 512
            # movingimg = movingimg.reshape(int(t), 512, 512) # z y x
            # movingimg = movingimg.transpose(2, 1, 0)
            # movingimg = movingimg.astype(np.float32)

            #   读取标签
            if labeldir is not None:
                label = nib.load(labelpath)
                label = label.get_fdata()
                label = label.astype(np.float32)

                #   同理
                # label = np.fromfile(file=labelpath, dtype=np.int16)
                # t = label.shape[0] / 512 / 512
                # label = label.reshape(int(t), 512, 512)
                # label = label.transpose(2, 1, 0)

            #   生成参数对象，后面用于承受txt的参数文件
            parameter_object = itk.ParameterObject.New()

            # 参数文件可上Elastix model Zoo下载， par044是用于血管CT配准的
            # parameter_object.AddParameterFile('ElastixParameters/Par0059_rigid.txt')
            # parameter_object.AddParameterFile('ElastixParameters/Par0059_bspline.txt')
            parameter_object.AddParameterFile('ElastixParameters/par044affine.txt')
            parameter_object.AddParameterFile('ElastixParameters/par044norigid.txt')
            # parameter_object.AddParameterFile('ElastixParameters/affine.txt')
            # parameter_object.AddParameterFile('ElastixParameters/Bspline.txt')

            # 对image进行变换, log参数可以选择控制台打印过程日志，一般很长所以除非报错找不到原因，不然就不开
            result_image, result_transform_parameters = itk.elastix_registration_method(fixedimg, movingimg, parameter_object=parameter_object, log_to_console=log)
            result_image = result_image.astype(np.float32)
            result_image = result_image.transpose(2, 1, 0)
            result_image = nib.Nifti1Image(result_image, affine)
            nib.save(result_image, saveimg_)

            # 对label进行变换
            if savelabeldir is not None:
                label = itk.GetImageFromArray(label)
                transformix_object = itk.TransformixFilter.New(label)
                transformix_object.SetTransformParameterObject(result_transform_parameters)
                transformix_object.UpdateLargestPossibleRegion()

                # 配准好的label保存
                result_moving_mask = transformix_object.GetOutput()
                result_moving_mask = itk.GetArrayFromImage(result_moving_mask)
                result_moving_mask = result_moving_mask.transpose(2, 1, 0)
                result_moving_mask[result_moving_mask >= 0.5] = 1
                result_moving_mask[result_moving_mask < 0.5] = 0
                result_moving_mask = result_moving_mask.astype(np.int16)
                result_moving_mask = nib.Nifti1Image(result_moving_mask, affine)
                nib.save(result_moving_mask, savelbl)

        except Exception as r:
            print(r)


def clip(fixedimgdir, savedir, savelabeldir=None):
    """
    裁剪函数，配准后一般会有区域无效（全黑），为了不影响后续模型训练，可以用该函数进行切除
    但是需要手动筛选出区域，并记录在excel中，表格格式为：
    name | start_slice | end_slice
    不需要表头，从第一行开始就直接记录，name为数据的文件名，start_slice和end_slice约束了一个保留范围，层数的计数以0为起始，len-1为终止
    """
    excel_list = pd.read_excel(io=r'need_clip.xlsx', header=None)
    for i in tqdm(range(len(excel_list))):
        name, start, end = str(excel_list[0][i]), int(excel_list[1][i]), int(excel_list[2][i])

        fixed_path = os.path.join(fixedimgdir, name+'.nii')
        affine_path = os.path.join(savedir, name+'.nii')
        if savelabeldir is not None:
            label_path = os.path.join(savelabeldir, name+'.nii')

        fixed_img = nib.load(fixed_path)
        fixed_affine = fixed_img.affine
        fixed_img = fixed_img.get_fdata().astype(np.float32)

        affine_img = nib.load(affine_path)
        affine_affine = affine_img.affine
        affine_img = affine_img.get_fdata().astype(np.float32)

        if savelabeldir is not None:
            label_img = nib.load(label_path)
            label_affine = label_img.affine
            label_img = label_img.get_fdata().astype(np.float32)

        nib.save(nib.Nifti1Image(fixed_img[:, :, start:end+1], fixed_affine), fixed_path)
        nib.save(nib.Nifti1Image(affine_img[:, :, start:end + 1], affine_affine), affine_path)
        if savelabeldir is not None:
            nib.save(nib.Nifti1Image(label_img[:, :, start:end + 1], label_affine), label_path)
        break


if __name__ == '__main__':
    #   待配准的图像路径
    moving_img_dir = r"xxx"
    #   待配准的图像的标签路径
    label_dir = r"xxx"

    #   目标的图像路径
    fixed_img_dir = r"xxx"

    #   配准好后的图像和标签存放路径
    save_dir = r"xxx"
    save_label_dir = r"xxx"

    #   调用配准函数
    reg(moving_img_dir, fixed_img_dir, label_dir, save_dir, save_label_dir)

    #   调用裁剪函数
    clip(fixed_img_dir, save_dir, save_label_dir)
